import numpy as np


class PendulumODE:
    MU = 3.9860044 * (10**14)
    AE = 6378136
    J2 = 0.0010826257

    def __init__(
        self,
        mu: float = MU,
        ae: int = AE,
        j2: float = J2,
    ):
        """PendulumODE init function. Sets init parameteres.

        Parameters
        ----------
        mu : float
            Gravity parameter
        ae : int
            The semi-axis of the Earth ellipsoid
        j2 : float
            J2 coefficient
        """
        self.mu = mu
        self.ae = ae
        self.j2 = j2

    def pendulum_np(self, t: float, x: list) -> np.ndarray:
        """Calculate coordinates and velocities with 2nd harmonic.

        Parameters
        ----------
        t : float
            Time
        x : list
            (1x6) initial coordinates and velocities vector (x, y, z, x', y', z')

        Returns
        -------
        np.ndarray
            (1x6) coordinates and velocities vector (x, y, z, x', y', z')
        """
        r = np.sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2)
        dx1 = x[3]
        dx2 = x[4]
        dx3 = x[5]
        t = (2 * (r**5)) * (5 * (x[2] ** 2) / (r**2) - 1)
        z = (2 * (r**5)) * (5 * (x[2] ** 2) / (r**2) - 3 * x[2])
        dx4 = (
            -self.mu * x[0] / (r**3)
            + 3 * self.mu * (self.ae**2) * self.j2 * x[0] / t
        )
        dx5 = (
            -self.mu * x[1] / (r**3)
            + 3 * self.mu * (self.ae**2) * self.j2 * x[1] / t
        )
        dx6 = -self.mu * x[2] / (r**3) + 3 * self.mu * (self.ae**2) * self.j2 / z
        return np.array([dx1, dx2, dx3, dx4, dx5, dx6])

    def pendulum(self, t: float, x: list) -> np.ndarray:
        """Calculate coordinates and velocities without 2nd harmonic.

        Parameters
        ----------
        t : float
            Time
        x : list
            (1x6) initial coordinates and velocities vector (x, y, z, x', y', z')

        Returns
        -------
        np.ndarray
            (1x6) coordinates and velocities vector (x, y, z, x', y', z')
        """
        r = np.sqrt(x[0] ** 2 + x[1] ** 2 + x[2] ** 2)
        dx1 = x[3]
        dx2 = x[4]
        dx3 = x[5]
        dx4 = -self.mu * x[0] / (r**3)
        dx5 = -self.mu * x[1] / (r**3)
        dx6 = -self.mu * x[2] / (r**3)
        return np.array([dx1, dx2, dx3, dx4, dx5, dx6])
